"""
讯飞星火大模型 WebSocket API 调用模块
官方文档: https://www.xfyun.cn/doc/spark/Web.html
使用 Spark Lite 模型（免费版本）
"""
import _thread as thread
import base64
import hashlib
import hmac
import json
import ssl
from datetime import datetime
from time import mktime
from urllib.parse import urlencode, urlparse
from wsgiref.handlers import format_date_time

import websocket  # 使用 websocket-client 库

# 全局变量存储答案
answer = ''


class Ws_Param(object):
    """WebSocket 参数类，用于生成鉴权 URL"""
    
    def __init__(self, APPID, APIKey, APISecret, Spark_url):
        """
        初始化 WebSocket 参数
        
        Args:
            APPID: 应用 ID
            APIKey: API Key
            APISecret: API Secret
            Spark_url: 星火大模型 WebSocket 地址
        """
        self.APPID = APPID
        self.APIKey = APIKey
        self.APISecret = APISecret
        self.host = urlparse(Spark_url).netloc
        self.path = urlparse(Spark_url).path
        self.Spark_url = Spark_url

    def create_url(self):
        """
        生成鉴权 URL
        
        Returns:
            str: 带鉴权参数的 WebSocket URL
        """
        # 生成 RFC1123 格式的时间戳
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))

        # 拼接签名字符串
        signature_origin = "host: " + self.host + "\n"
        signature_origin += "date: " + date + "\n"
        signature_origin += "GET " + self.path + " HTTP/1.1"

        # 使用 HMAC-SHA256 加密
        signature_sha = hmac.new(
            self.APISecret.encode('utf-8'),
            signature_origin.encode('utf-8'),
            digestmod=hashlib.sha256
        ).digest()

        signature_sha_base64 = base64.b64encode(signature_sha).decode(encoding='utf-8')

        # 生成 authorization 字段
        authorization_origin = (
            f'api_key="{self.APIKey}", '
            f'algorithm="hmac-sha256", '
            f'headers="host date request-line", '
            f'signature="{signature_sha_base64}"'
        )

        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode(encoding='utf-8')

        # 组合鉴权参数
        v = {
            "authorization": authorization,
            "date": date,
            "host": self.host
        }
        
        # 拼接完整 URL
        url = self.Spark_url + '?' + urlencode(v)
        return url


def on_error(ws, error):
    """WebSocket 错误处理"""
    print(f"### WebSocket Error: {error}")


def on_close(ws, close_status_code, close_msg):
    """WebSocket 关闭处理"""
    pass


def on_open(ws):
    """WebSocket 连接建立处理"""
    thread.start_new_thread(run, (ws,))


def run(ws, *args):
    """发送请求数据"""
    data = json.dumps(gen_params(
        appid=ws.appid,
        domain=ws.domain,
        query=ws.query
    ))
    ws.send(data)


def on_message(ws, message):
    """
    WebSocket 消息处理
    
    Args:
        ws: WebSocket 对象
        message: 接收到的消息
    """
    data = json.loads(message)
    code = data['header']['code']
    
    if code != 0:
        print(f'请求错误: {code}, {data}')
        ws.close()
    else:
        choices = data["payload"]["choices"]
        status = choices["status"]
        content = choices["text"][0]["content"]
        
        # 累积答案内容
        global answer
        answer += content
        
        # status=2 表示响应结束
        if status == 2:
            ws.close()


def gen_params(appid, domain, query):
    """
    生成请求参数
    
    Args:
        appid: 应用 ID
        domain: 模型领域
        query: 用户问题
        
    Returns:
        dict: 请求参数字典
    """
    data = {
        "header": {
            "app_id": appid,
            "uid": "1234",
        },
        "parameter": {
            "chat": {
                "domain": domain,
                "temperature": 0.5,
                "max_tokens": 4096,
                "auditing": "default",
            }
        },
        "payload": {
            "message": {
                "text": [
                    {
                        "role": "system",
                        "content": "在选项中找出问题的正确答案，请保证信息的准确，并在回答时保持简洁，不需要任何其他分析反馈,只说正确选项即可。"
                    },
                    {
                        "role": "user",
                        "content": query
                    }
                ]
            }
        }
    }
    return data


def get_answer(appid, api_secret, api_key, question):
    """
    调用讯飞星火大模型获取答案（WebSocket 方式）
    使用 Spark Lite 模型（免费版本）
    
    Args:
        appid: 应用 ID
        api_secret: API Secret
        api_key: API Key
        question: 问题内容
    """
    global answer
    answer = '│ XF：'
    
    # Spark Max 模型配置
    Spark_url = 'wss://spark-api.xf-yun.com/v3.5/chat'
    domain = 'generalv3.5'

    # Spark Lite 模型配置（免费版本）
    # Spark_url = 'wss://spark-api.xf-yun.com/v1.1/chat'
    # domain = 'lite'
    
    # 创建 WebSocket 参数
    wsParam = Ws_Param(appid, api_key, api_secret, Spark_url)
    websocket.enableTrace(False)
    wsUrl = wsParam.create_url()
    
    # 创建 WebSocket 连接
    ws = websocket.WebSocketApp(
        wsUrl,
        on_message=on_message,
        on_error=on_error,
        on_close=on_close,
        on_open=on_open
    )
    ws.appid = appid
    ws.query = question
    ws.domain = domain
    
    # 运行 WebSocket 连接
    ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})
    
    # 打印完整答案
    print(answer)
